﻿#include "VulkanSpotLights.h"
#include "Render/Texture/GPURenderTarget.h"
#include "Render/Scene/VulkanScene.h"
#include "Render/Scene/DefferRenderAttachments.h"
#include "Render/Widget/VulkanViewportWidget.h"

void CVulkanSpotLights::PreProcess(const TVector<SPreRenderSpotLight>& InSpotLights, CVulkanViewportWidget * InViewport)
{
    //剔除与视椎体不相交的聚光灯
    SpotLights.Clear();

    for (const auto& SpotLight : InSpotLights)
    {
        //检查参数是否合法
        if (SpotLight.Radius <= 0.0f || SpotLight.Intensity <= 0.0f)
        {
            continue;
        }

        if(SpotLight.ConeAngle<=0.0f || SpotLight.ConeAngle>=180.0f)
        {
            continue;
        }

        if(SpotLight.Color.R<=0.0f && SpotLight.Color.G<=0.0f && SpotLight.Color.B<=0.0f)
        {
            //黑色光源
            continue;
        }

        //检查是否在视椎体内
        if (!InViewport->FastIsSphereIntersectFrustum(SpotLight.Position, SpotLight.Radius))
        {
            continue;
        }


        SpotLights.Add(SpotLight);
    }

    //统计描述符数量
    SVulkanDescriptorPoolSize PoolSize;
    PoolSize+=CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapDescriptorSetLayout.DescriptorNums*SpotLights.Num();
    PoolSize.SetNum+=SpotLights.Num();
    CVulkanContext::Get()->GetAutoResetDescriptorPool()->GenerallPool.PoolSize+=PoolSize;

    SVulkanDescriptorPoolSize PoolSize2;
    PoolSize2+=CVulkanContext::Get()->GetSceneGlobal()->SpotLightDefferShadingDescriptorSetLayout.DescriptorNums*SpotLights.Num();
    PoolSize2.SetNum+=SpotLights.Num();
    CVulkanContext::Get()->GetAutoResetDescriptorPool()->GenerallPool.PoolSize+=PoolSize2;

    //绑定数据
    RenderLights.Resize(SpotLights.Num());
    for(int32_t Index=0;Index<SpotLights.Num();Index++)
    {
        RenderLights[Index].SetSpotLight(SpotLights[Index]);
        RenderLights[Index].LightIndex=Index;
    }
}

void CVulkanSpotLights::RecordShadowMapCommand(vk::CommandBuffer InCmdBuffers, CVulkanScene *InScene)
{
    for (auto& RenderLight : RenderLights)
    {
        RenderLight.RecordShadowMapCommand(InCmdBuffers, InScene);
    }
}

void CVulkanSpotLights::RecordDefferShadingCommand(vk::CommandBuffer InCmdBuffers, CVulkanScene *InScene, const SSceneParameters &InParameters, CDefferRenderAttachments *Attachments)
{
    InCmdBuffers.bindPipeline(vk::PipelineBindPoint::eGraphics, CVulkanContext::Get()->GetSceneGlobal()->SpotLightDefferShadingPipeline);

    for (auto& RenderLight : RenderLights)
    {
        RenderLight.RecordDefferShadingCommand(InCmdBuffers, InScene, InParameters, Attachments);
    }
}

bool CVulkanSpotLights::IsInLightRange(int64_t InLightIndex, const SVec3 &InSpherePosition, float InSphereRadius)
{
    assert(InLightIndex >= 0 && InLightIndex < RenderLights.Num());

    return (SpotLights[InLightIndex].Position-InSpherePosition).GetLength() < +InSphereRadius+SpotLights[InLightIndex].Radius;
}

CVulkanSpotLight::CVulkanSpotLight()
{
    //生成缓冲区
    ShadowMapBuffer = std::make_shared<TVulkanUniformBuffer<CSceneGlobal::SSpotLightShadowMapBuffer>>();

    //创建渲染目标
    Resize(CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapExtent.width
    , CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapExtent.height);
}

CVulkanSpotLight::~CVulkanSpotLight()
{
    if (ShadowMapFrameBuffer)
    {
        CVulkanContext::Get()->GetDevice().destroyFramebuffer(ShadowMapFrameBuffer);
    }
}

void CVulkanSpotLight::SetSpotLight(const SPreRenderSpotLight &InSpotLight)
{
    RenderData = InSpotLight;
    
    SMatrix4x4 ProjectionMatrix = SMatrix4x4::Perspective(
        RenderData.ConeAngle
        , 1.0f
        , 0.1f
        , RenderData.Radius
    );

    SMatrix4x4 ViewMatrix = SMatrix4x4::LookAt(
        RenderData.Position
        , RenderData.Rotation.GetForwardDirection()
        , RenderData.Rotation.GetUpDirection()
    );

    ProjectionViewMatrix = ProjectionMatrix * ViewMatrix;

    //更新缓冲区
    CSceneGlobal::SSpotLightShadowMapBuffer ShadowMapBufferContent;
    ProjectionViewMatrix.ToFloatArray(ShadowMapBufferContent.ProjectionViewMatrix);
    RenderData.Position.ToFloatArray(ShadowMapBufferContent.LightPosition);
    ShadowMapBuffer->WriteData(ShadowMapBufferContent);
}

void CVulkanSpotLight::RecordShadowMapCommand(vk::CommandBuffer InCmdBuffers, CVulkanScene *InScene)
{
    //重新生成描述符,由于描述符池自动重置，所以每帧都要重新生成
    vk::DescriptorSetAllocateInfo DescriptorSetAllocateInfo = vk::DescriptorSetAllocateInfo()
            .setDescriptorPool(CVulkanContext::Get()->GetAutoResetDescriptorPool()->GenerallPool.Pool)
            .setDescriptorSetCount(1)
            .setPSetLayouts(&CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapDescriptorSetLayout.Layout);

    ShadowMapDescriptorSet = CVulkanContext::Get()->GetDevice().allocateDescriptorSets(DescriptorSetAllocateInfo)[0];

    vk::DescriptorBufferInfo DescriptorBufferInfo = vk::DescriptorBufferInfo()
        .setBuffer(ShadowMapBuffer->GetBuffer())
        .setOffset(0)
        .setRange(ShadowMapBuffer->GetSize());

    vk::WriteDescriptorSet WriteDescriptorSet = vk::WriteDescriptorSet()
        .setDstSet(ShadowMapDescriptorSet)
        .setDstBinding(0)
        .setDescriptorType(vk::DescriptorType::eUniformBuffer)
        .setDescriptorCount(1)
        .setPBufferInfo(&DescriptorBufferInfo);

    CVulkanContext::Get()->GetDevice().updateDescriptorSets(WriteDescriptorSet, nullptr);    

    //记录命令
    TVector<vk::ClearValue> ClearValues =
    {
        vk::ClearColorValue(std::array<float,4>{-1.f /*无符号的最大值*/,0.0f,0.0f,0.0f}),
        vk::ClearDepthStencilValue(1.0f, 0)
    };   

    auto SpotLightRenderPassBeginInfo=vk::RenderPassBeginInfo()
            .setRenderPass(CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapRenderPass)
            .setFramebuffer(ShadowMapFrameBuffer)
            .setRenderArea(vk::Rect2D(vk::Offset2D(0,0),vk::Extent2D(ShadowMap->GetSize().width,ShadowMap->GetSize().height)))
            .setClearValues(ClearValues);

    InCmdBuffers.beginRenderPass(SpotLightRenderPassBeginInfo, vk::SubpassContents::eInline);     

    InCmdBuffers.bindPipeline(vk::PipelineBindPoint::eGraphics, CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapPipeline);

    //渲染所有的Mesh
    for (const auto& MeshInstanceIndex : InScene->GetMeshInstancePainter()->SpotLightShadowMapMeshMap[LightIndex])
    {
        auto& MeshInstance =InScene->GetMeshInstancePainter()->MeshInstances[MeshInstanceIndex];

        //绑定描述符集
        InCmdBuffers.bindDescriptorSets(
            vk::PipelineBindPoint::eGraphics
            ,CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapPipelineLayout
            ,0
            ,{ ShadowMapDescriptorSet, MeshInstance.GetModelUniformDescriptorSet() }
            ,{}
        );

        
        MeshInstance.RecordDrawVertexCommands(InCmdBuffers);
                     
    }

    InCmdBuffers.endRenderPass();

    //将阴影贴图布局转换为eShaderReadOnlyOptimal
    vk::ImageMemoryBarrier ImageMemoryBarrier = vk::ImageMemoryBarrier()
        .setOldLayout(vk::ImageLayout::eColorAttachmentOptimal)
        .setNewLayout(vk::ImageLayout::eShaderReadOnlyOptimal)
        .setSrcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
        .setDstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
        .setImage(ShadowMap->GetImage())
        .setSubresourceRange(vk::ImageSubresourceRange(vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1))
        .setSrcAccessMask(vk::AccessFlagBits::eColorAttachmentWrite)
        .setDstAccessMask(vk::AccessFlagBits::eShaderRead);
    
    InCmdBuffers.pipelineBarrier(
        vk::PipelineStageFlagBits::eColorAttachmentOutput,
        vk::PipelineStageFlagBits::eFragmentShader,
        vk::DependencyFlagBits::eByRegion,
        nullptr,
        nullptr,
        ImageMemoryBarrier
    );    
}

void CVulkanSpotLight::RecordDefferShadingCommand(vk::CommandBuffer InCmdBuffers, CVulkanScene *InScene, const SSceneParameters &InParameters, CDefferRenderAttachments *Attachments)
{
    //创建描述符
    SVulkanDescriptorLayout DescriptorLayout=
        CVulkanContext::Get()->GetSceneGlobal()->SpotLightDefferShadingDescriptorSetLayout;

    vk::DescriptorSetAllocateInfo DescriptorSetAllocateInfo = vk::DescriptorSetAllocateInfo()
            .setDescriptorPool(CVulkanContext::Get()->GetAutoResetDescriptorPool()->GenerallPool.Pool)
            .setDescriptorSetCount(1)
            .setPSetLayouts(&DescriptorLayout.Layout);

    DefferShadingDescriptorSet =CVulkanContext::Get()->GetDevice().allocateDescriptorSets(DescriptorSetAllocateInfo)[0];

    //update descriptor set

    vk::DescriptorImageInfo DescriptorImageInfo= vk::DescriptorImageInfo()
        .setImageLayout(vk::ImageLayout::eShaderReadOnlyOptimal)
        .setImageView(ShadowMap->GetImageView())
        .setSampler(CVulkanContext::Get()->GetSceneGlobal()->ShadowMapSampler)
        ;

    TVector<vk::WriteDescriptorSet> WriteDescriptorSets=
    {
            vk::WriteDescriptorSet()
                .setDstSet(DefferShadingDescriptorSet)
                .setDstBinding(0)
                .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
                .setDescriptorCount(1)
                .setPImageInfo(&DescriptorImageInfo)
    };

    CVulkanContext::Get()->GetDevice().updateDescriptorSets(WriteDescriptorSets, nullptr);    

    //绘制
    InCmdBuffers.bindDescriptorSets(vk::PipelineBindPoint::eGraphics
        ,CVulkanContext::Get()->GetSceneGlobal()->SpotLightDefferShadingPipelineLayout
        , 0
        ,{Attachments->DescriptorSet, DefferShadingDescriptorSet}
        ,{}
    );

    //push constants
    SSpotLightDefferShadingFragmentPushConstant PushConstant;

    (RenderData.Color*RenderData.Intensity).ToFloatArray(PushConstant.ColorWithIntensity);
    RenderData.Position.ToFloatArray(PushConstant.Position);
    ProjectionViewMatrix.ToFloatArray(PushConstant.LightVP);
    PushConstant.ViewPos[0] = InParameters.CameraPosition[0];
    PushConstant.ViewPos[1] = InParameters.CameraPosition[1];
    PushConstant.ViewPos[2] = InParameters.CameraPosition[2];
    PushConstant.ShadowBias =RenderData.ShadowBias;

    InCmdBuffers.pushConstants(
            CVulkanContext::Get()->GetSceneGlobal()->SpotLightDefferShadingPipelineLayout
            , vk::ShaderStageFlagBits::eFragment
            , 0
            , sizeof(SSpotLightDefferShadingFragmentPushConstant)
            , &PushConstant
    );

    InCmdBuffers.draw(3, 1, 0, 0);


}

void CVulkanSpotLight::Resize(uint32_t InWidth, uint32_t InHeight)
{
    if (Width == InWidth && Height == InHeight)
    {
        return;
    }

    //清除旧资源
    if (ShadowMapFrameBuffer)
    {
        CVulkanContext::Get()->GetDevice().destroyFramebuffer(ShadowMapFrameBuffer);
    }

    Width = InWidth;
    Height = InHeight;

    //创建渲染目标
    ShadowMap = std::make_shared<CGPURenderTarget>();
    ShadowMap->Initialize(
        vk::Extent3D(Width, Height, 1)
        , CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapImageFormat
        , vk::ImageUsageFlagBits::eColorAttachment | vk::ImageUsageFlagBits::eSampled
    );

    Depth = std::make_shared<CGPURenderTarget>();
    Depth->Initialize(
        vk::Extent3D(Width, Height, 1)
        , vk::Format::eD16Unorm
        , vk::ImageUsageFlagBits::eDepthStencilAttachment
    );

    TVector<vk::ImageView> Attachments =
    {
        ShadowMap->GetImageView(),
        Depth->GetImageView()
    };

    //创建帧缓冲
    vk::FramebufferCreateInfo FrameBufferCreateInfo = vk::FramebufferCreateInfo()
        .setRenderPass(CVulkanContext::Get()->GetSceneGlobal()->SpotLightShadowMapRenderPass)
        .setAttachmentCount(Attachments.Num())
        .setPAttachments(Attachments.Data())
        .setWidth(Width)
        .setHeight(Height)
        .setLayers(1);

    ShadowMapFrameBuffer = CVulkanContext::Get()->GetDevice().createFramebuffer(FrameBufferCreateInfo);

}
